using System;
using System.Collections.Generic;
using System.Compiler;
using System.Diagnostics.Contracts;
using System.Linq;
using System.Threading.Tasks;
using Microsoft.Contracts.Foxtrot.Utils;

namespace Microsoft.Contracts.Foxtrot
{
    /// <summary>
    /// Visitor that creates special closure type for async postconditions.
    /// </summary>
    /// <remarks>
    /// 
    /// Current class generate AsyncClosure with CheckMethod and CheckException methods.
    /// 
    /// Following transformation are applied to the original async method:
    /// 
    /// // This is oroginal task, generated by the AsyncTaskMethodBuilder
    /// var originalTask = t_builder.Task;
    /// 
    /// var closure = new AsyncClosure();
    /// var task2 = originalTask.ContinueWith(closure.CheckPost).Unwrap();
    /// return task2;
    /// 
    /// There are 2 cases:
    /// 1) Task has no return value.
    /// In this case only EnsuresOnThrow could be used, and we emit:
    /// Task CheckMethod(Task t)
    /// {
    ///   if (t.Status == TaskStatus.Faulted)
    ///   {
    ///     // CheckException will throw if EnsuresOnThrow is not held
    ///     CheckException(t.Exception);
    ///   }
    /// 
    ///   return t;
    /// }
    /// 
    /// 2) Task(T) reutrns a T value.
    /// In this case both EnsuresOnThrow and Contract.Ensures(Contract.Result) could be used.
    /// We emit:
    /// 
    /// Task&lt;int> CheckMethod(Task&lt;int> t)
    /// {
    ///   if (t.Status == TaskStatus.Faulted)
    ///   {
    ///     // CheckException will throw if EnsuresOnThrow is not held
    ///     CheckException(t.Exception);
    ///   }
    ///   
    ///   if (t.Status == TaskStatus.RanToCompletion)
    ///   {
    ///      // Check ensures
    ///   }
    /// }
    /// </remarks>
    internal class EmitAsyncClosure : StandardVisitor
    {
        /// <summary>
        /// Class that maps generic arguments of the enclosed class/method to the generic arguments of the closure.
        /// </summary>
        /// <remarks>
        /// The problem.
        /// Original implementation of the Code Contract didn't support async postconditions in generics.
        /// Here is why:
        /// Suppose we have following function (in class <c>Foo</c>:
        /// <code><![CDATA[
        /// public static Task<T> FooAsync() where T: class
        /// {
        ///     Contract.Ensures(Contract.Result<T>() != null);
        /// }
        /// ]]></code>
        /// In this case, ccrewrite will generate async closure class called <c>Foo.AsyncContractClosure_0&lt;T&gt;</c>
        /// with following structure:
        /// <code><![CDATA[
        /// [CompilerGenerated]
        /// private class <Foo>AsyncContractClosure_0<T> where T : class
        /// {
        /// 	public Task<T> CheckPost(Task<T> task)
        /// 	{
        /// 		TaskStatus status = task.Status;
        /// 		if (status == TaskStatus.RanToCompletion)
        /// 		{
        /// 			RewriterMethods.Ensures(task.Result != null, null, "Contract.Result<T>() != null");
        /// 		}
        /// 		return task;
        /// 	}
        /// }
        /// ]]>
        /// </code>
        /// The code looks good, but the IL could be invalid (without the trick that this class provides).
        /// Postcondition of the method in our case is declared in the generic method (in <code>FooAsync</code>)
        /// but ccrewrite moves it into non-generic method (<code>CheckPost</code>) of the generic class (closure).
        /// 
        /// But on IL level there is different instructions for referencing method generic arguments and type generic arguments.
        /// 
        /// After changing <code>Contract.Result</code> to <code>task.Result</code> and moving postcondition to <code>CheckPost</code>
        /// method, following IL would be generated:
        /// 
        /// <code> <![CDATA[
        /// IL_0011: call instance !0 class [mscorlib]System.Threading.Tasks.Task`1<!T>::get_Result()
        /// IL_0016: box !!0 // <-- here is our problem!
        /// ]]>
        /// </code>
        ///
        /// This means that method <code>CheckPost</code> would contains a reference to generic method argument of the
        /// original method.
        /// 
        /// The goal of this class is to store a mapping between enclosing generic types and closure generic types.
        /// </remarks>
        private class GenericTypeMapper
        {
            class TypeNodePair
            {
                public TypeNodePair(TypeNode enclosingGenericType, TypeNode closureGenericType)
                {
                    EnclosingGenericType = enclosingGenericType;
                    ClosureGenericType = closureGenericType;
                }

                public TypeNode EnclosingGenericType { get; private set; }
                public TypeNode ClosureGenericType { get; private set; }
            }

            // Mapping between enclosing generic type and closure generic type.
            // This is a simple list not a dictionary, because number of generic arguments is very small.
            // So linear complexity will not harm performance.
            readonly List<TypeNodePair> typeParametersMapping = new List<TypeNodePair>();

            public bool IsEmpty { get { return typeParametersMapping.Count == 0; } }

            public void AddMapping(TypeNode enclosingGenericType, TypeNode closureGenericType)
            {
                typeParametersMapping.Add(new TypeNodePair(enclosingGenericType, closureGenericType));
            }

            /// <summary>
            /// Returns associated generic type of the closure class by enclosing generic type (for instance, by
            /// generic type of the enclosing generic method that uses current closure).
            /// </summary>
            /// <remarks>
            /// Function returns the same argument if the matching argument does not exists.
            /// </remarks>
            public TypeNode GetClosureTypeParameterByEnclosingTypeParameter(TypeNode enclosingType)
            {
                if (enclosingType == null)
                {
                    return null;
                }

                var gen = enclosingType;
                if (gen.ConsolidatedTemplateParameters != null && gen.ConsolidatedTemplateParameters.Count != 0)
                {
                    gen = gen.ConsolidatedTemplateParameters[0];
                }

                var candidate = typeParametersMapping.FirstOrDefault(t => t.EnclosingGenericType == enclosingType);
                return candidate != null ? candidate.ClosureGenericType : enclosingType;
            }
            
            /// <summary>
            /// Returns associated generic type of the closure class by enclosing generic type (for instance, by
            /// generic type of the enclosing generic method that uses current closure).
            /// </summary>
            /// <remarks>
            /// Function returns the same argument if the matching argument does not exists.
            /// </remarks>
            public TypeNode GetEnclosingTypeParameterByClosureTypeParameter(TypeNode closureType)
            {
                if (closureType == null)
                {
                    return null;
                }

                var candidate = typeParametersMapping.FirstOrDefault(t => t.ClosureGenericType == closureType);
                return candidate != null ? candidate.EnclosingGenericType : closureType;
            }
        }

        // This assembly should be in this class but not in the SystemTypes from System.CompilerCC.
        // Moving this type there will lead to test failures and assembly resolution errors.
        private static readonly AssemblyNode/*!*/ SystemCoreAssembly = SystemTypes.GetSystemCoreAssembly(false, true);
        
        private static TypeNode TaskExtensionsTypeNode = HelperMethods.FindType(
            SystemCoreAssembly,
            Identifier.For("System.Threading.Tasks"),
            Identifier.For("TaskExtensions"));

        private static readonly Identifier CheckExceptionMethodId = Identifier.For("CheckException");
        private static readonly Identifier CheckMethodId = Identifier.For("CheckPost");

        private readonly Cache<TypeNode> aggregateExceptionType;
        private readonly Cache<TypeNode> func2Type;

        private readonly Dictionary<Local, MemberBinding> closureLocals = new Dictionary<Local, MemberBinding>();
        private readonly List<SourceContext> contractResultCapturedInStaticContext = new List<SourceContext>();

        private readonly Rewriter rewriter;
        private readonly TypeNode declaringType;
        private readonly Class closureClass;
        private readonly Class closureClassInstance;
        private readonly Specializer /*?*/ forwarder;
        private readonly Local closureLocal;

        // Fields for the CheckMethod generation
        private Method checkPostMethod;
        private StatementList checkPostBody;
        // Holds a copy of CheckMethod argument
        private Local originalResultLocal;
        private Parameter checkMethodTaskParameter;
        private readonly TypeNode checkMethodTaskType;

        private readonly GenericTypeMapper genericTypeMapper = new GenericTypeMapper();

        public EmitAsyncClosure(Method from, Rewriter rewriter)
        {
            Contract.Requires(from != null);
            Contract.Requires(from.DeclaringType != null);
            Contract.Requires(rewriter != null);

            if (TaskExtensionsTypeNode == null)
            {
                throw new InvalidOperationException(
                    "Can't generate async closure because System.Threading.Tasks.TaskExceptions class is unavailable.");
            }

            this.rewriter = rewriter;
            this.declaringType = from.DeclaringType;

            var closureName = HelperMethods.NextUnusedMemberName(declaringType, "<" + from.Name.Name + ">AsyncContractClosure");

            this.closureClass = new Class(
                declaringModule: declaringType.DeclaringModule,
                declaringType: declaringType,
                attributes: null,
                flags: TypeFlags.NestedPrivate,
                Namespace: null,
                name: Identifier.For(closureName),
                baseClass: SystemTypes.Object,
                interfaces: null,
                members: null);

            declaringType.Members.Add(this.closureClass);
            RewriteHelper.TryAddCompilerGeneratedAttribute(this.closureClass);

            var taskType = from.ReturnType;

            this.aggregateExceptionType = new Cache<TypeNode>(() =>
                HelperMethods.FindType(rewriter.AssemblyBeingRewritten, StandardIds.System,
                    Identifier.For("AggregateException")));

            this.func2Type = new Cache<TypeNode>(() =>
                HelperMethods.FindType(SystemTypes.SystemAssembly, StandardIds.System, Identifier.For("Func`2")));

            // Should distinguish between generic enclosing method and non-generic method in enclosing type.
            // In both cases generated closure should be generic.
            var enclosingTemplateParameters = GetGenericTypesFrom(from);

            if (!enclosingTemplateParameters.IsNullOrEmpty())
            {
                this.closureClass.TemplateParameters = CreateTemplateParameters(closureClass, enclosingTemplateParameters, declaringType);

                this.closureClass.IsGeneric = true;
                this.closureClass.EnsureMangledName();

                this.forwarder = new Specializer(
                    targetModule: this.declaringType.DeclaringModule,
                    pars: enclosingTemplateParameters,
                    args: this.closureClass.TemplateParameters);

                this.forwarder.VisitTypeParameterList(this.closureClass.TemplateParameters);

                taskType = this.forwarder.VisitTypeReference(taskType);

                for (int i = 0; i < enclosingTemplateParameters.Count; i++)
                {
                    this.genericTypeMapper.AddMapping(enclosingTemplateParameters[i], closureClass.TemplateParameters[i]);
                }
            }
            else
            {
                this.closureClassInstance = this.closureClass;
            }

            this.checkMethodTaskType = taskType;

            // Emiting CheckPost method declaration
            EmitCheckPostMethodCore(checkMethodTaskType);

            // Generate closure constructor.
            // Constructor should be generated AFTER visiting type parameters in
            // the previous block of code. Otherwise current class would not have
            // appropriate number of generic arguments!
            var ctor = CreateConstructor(closureClass);
            closureClass.Members.Add(ctor);

            // Now that we added the ctor and the check method, let's instantiate the closure class if necessary
            if (this.closureClassInstance == null)
            {
                var consArgs = new TypeNodeList();
                var args = new TypeNodeList();

                var parentCount = this.closureClass.DeclaringType.ConsolidatedTemplateParameters == null
                    ? 0
                    : this.closureClass.DeclaringType.ConsolidatedTemplateParameters.Count;

                for (int i = 0; i < parentCount; i++)
                {
                    consArgs.Add(this.closureClass.DeclaringType.ConsolidatedTemplateParameters[i]);
                }

                if (!enclosingTemplateParameters.IsNullOrEmpty())
                {
                    for (int i = 0; i < enclosingTemplateParameters.Count; i++)
                    {
                        consArgs.Add(enclosingTemplateParameters[i]);
                        args.Add(enclosingTemplateParameters[i]);
                    } 
                }

                this.closureClassInstance =
                    (Class)
                        this.closureClass.GetConsolidatedTemplateInstance(this.rewriter.AssemblyBeingRewritten,
                            closureClass.DeclaringType, closureClass.DeclaringType, args, consArgs);
            }

            // create closure initializer for context method
            this.closureLocal = new Local(this.ClosureClass);
            this.ClosureInitializer = new Block(new StatementList());

            // Generate constructor call that initializes closure instance
            this.ClosureInitializer.Statements.Add(
                new AssignmentStatement(
                    this.closureLocal,
                    new Construct(new MemberBinding(null, this.Ctor), new ExpressionList())));
        }

        /// <summary>
        /// Add postconditions for the task-based methods.
        /// </summary>
        /// <remarks>
        /// Method inserts all required logic to the <paramref name="returnBlock"/> calling
        /// ContinueWith method on the <paramref name="taskBasedResult"/>.
        /// </remarks>
        public void AddAsyncPostconditions(List<Ensures> asyncPostconditions, Block returnBlock, Local taskBasedResult)
        {
            Contract.Requires(asyncPostconditions != null);
            Contract.Requires(returnBlock != null);
            Contract.Requires(taskBasedResult != null);
            Contract.Requires(asyncPostconditions.Count > 0);

            // Async postconditions are impelemented using custom closure class
            // with CheckPost method that checks postconditions when the task
            // is finished.

            // Add Async postconditions to the AsyncClosure
            AddAsyncPost(asyncPostconditions);

            // Add task.ContinueWith().Unwrap(); method call to returnBlock
            AddContinueWithMethodToReturnBlock(returnBlock, taskBasedResult);

            ChangeThisReferencesToClosureLocals();
        }

        private void ChangeThisReferencesToClosureLocals()
        {
            var fieldRewriter = new FieldRewriter(this);
            fieldRewriter.Visit(this.checkPostMethod);
        }

        /// <summary>
        /// Returns a list of source spans where non-capturing lambdas were used.
        /// </summary>
        public IList<SourceContext> ContractResultCapturedInStaticContext
        {
            get { return contractResultCapturedInStaticContext; }
        }

        /// <summary>
        /// Instance used in calling method context
        /// </summary>
        public Class ClosureClass
        {
            get { return this.closureClassInstance; }
        }

        /// <summary>
        /// Local instance of the async closure class
        /// </summary>
        public Local ClosureLocal { get { return this.closureLocal; } }
        
        /// <summary>
        /// Block of code, responsible for closure instance initialization
        /// </summary>
        public Block ClosureInitializer { get; private set; }

        private InstanceInitializer Ctor
        {
            get { return (InstanceInitializer)this.closureClassInstance.GetMembersNamed(StandardIds.Ctor)[0]; }
        }

        private static TypeNodeList GetGenericTypesFrom(Method method)
        {
            if (method.IsGeneric)
            {
                return method.TemplateParameters;
            }

            if (method.DeclaringType.IsGeneric)
            {
                return GetFirstNonEmptyGenericListWalkingUpDeclaringTypes(method.DeclaringType);
            }

            return null;
        }

        private static TypeNodeList GetFirstNonEmptyGenericListWalkingUpDeclaringTypes(TypeNode node)
        {
            if (node == null)
            {
                return null;
            }

            if (node.TemplateParameters != null && node.TemplateParameters.Count != 0)
            {
                return node.TemplateParameters;
            }

            return GetFirstNonEmptyGenericListWalkingUpDeclaringTypes(node.DeclaringType);
        }

        [Pure]
        private static TypeNodeList CreateTemplateParameters(Class closureClass, TypeNodeList inputTemplateParameters, TypeNode declaringType)
        {
            Contract.Requires(closureClass != null);
            Contract.Requires(inputTemplateParameters != null);
            Contract.Requires(declaringType != null);

            var dup = new Duplicator(declaringType.DeclaringModule, declaringType);

            var templateParameters = new TypeNodeList();

            var parentCount = declaringType.ConsolidatedTemplateParameters.CountOrDefault();

            for (int i = 0; i < inputTemplateParameters.Count; i++)
            {
                var tp = HelperMethods.NewEqualTypeParameter(
                    dup, (ITypeParameter)inputTemplateParameters[i],
                    closureClass, parentCount + i);

                templateParameters.Add(tp);
            }

            return templateParameters;
        }

        private void AddContinueWithMethodToReturnBlock(Block returnBlock, Local taskBasedResult)
        {
            Contract.Requires(returnBlock != null);
            Contract.Requires(taskBasedResult != null);

            var taskType = taskBasedResult.Type;

            // To find appropriate ContinueWith method task type should be unwrapped
            var taskTemplate = HelperMethods.Unspecialize(taskType);

            var continueWithMethodLocal = GetContinueWithMethod(closureClass, taskTemplate, taskType);

            // TODO: not sure that this is possible situation when continueWith method is null. 
            // Maybe Contract.Assert(continueWithMethod != null) should be used instead!

            if (continueWithMethodLocal != null)
            {
                // We need to create delegate instance that should be passed to ContinueWith method
                var funcType = continueWithMethodLocal.Parameters[0].Type;
                var funcCtor = funcType.GetConstructor(SystemTypes.Object, SystemTypes.IntPtr);

                Contract.Assume(funcCtor != null);

                var funcLocal = new Local(funcCtor.DeclaringType);

                // Creating a method pointer to the AsyncClosure.CheckMethod
                // In this case we can't use checkMethod field.
                // Getting CheckMethod from clsoureClassInstance will provide correct (potentially updated)
                // generic arguments for enclosing type.
                var checkMethodFromClosureInstance = (Method) closureClassInstance.GetMembersNamed(CheckMethodId)[0];
                Contract.Assume(checkMethodFromClosureInstance != null);

                var ldftn = new UnaryExpression(
                    new MemberBinding(null, checkMethodFromClosureInstance),
                    NodeType.Ldftn,
                    CoreSystemTypes.IntPtr);

                // Creating delegate that would be used as a continuation for original task
                returnBlock.Statements.Add(
                    new AssignmentStatement(funcLocal,
                        new Construct(new MemberBinding(null, funcCtor),
                            new ExpressionList(closureLocal, ldftn))));

                // Wrapping continuation into TaskExtensions.Unwrap method
                // (this helps to preserve original exception and original result of the task,
                // but allows to throw postconditions violations).

                // Generating: result.ContinueWith(closure.CheckPost);
                var taskContinuationOption = new Literal(TaskContinuationOptions.ExecuteSynchronously);
                var continueWithCall =
                    new MethodCall(
                        new MemberBinding(taskBasedResult, continueWithMethodLocal),
                        new ExpressionList(funcLocal, taskContinuationOption));

                // Generating: TaskExtensions.Unwrap(result.ContinueWith(...))
                var unwrapMethod = GetUnwrapMethod(checkMethodTaskType);
                var unwrapCall =
                    new MethodCall(
                        new MemberBinding(null, unwrapMethod), new ExpressionList(continueWithCall));

                // Generating: result = Unwrap(...);
                var resultAssignment =
                    new AssignmentStatement(taskBasedResult, unwrapCall);
                returnBlock.Statements.Add(resultAssignment);
            }
        }

        /// <summary>
        /// Method generates core part of the CheckMethod
        /// </summary>
        private void EmitCheckPostMethodCore(TypeNode taskType)
        {
            Contract.Requires(taskType != null);

            this.checkMethodTaskParameter = new Parameter(Identifier.For("task"), taskType);

            // TODO ST: can I switch to new Local(taskType.Type)?!? In this case this initialization
            // could be moved outside this method
            this.originalResultLocal = new Local(new Identifier("taskLocal"), checkMethodTaskParameter.Type);

            // Generate: public Task<T> CheckPost(Task<T> task) where T is taskType or
            // public Task CheckPost(Task task) for non-generic task.
            checkPostMethod = new Method(
                declaringType: this.closureClass,
                attributes: null,
                name: CheckMethodId,
                parameters: new ParameterList(checkMethodTaskParameter),
                // was: taskType.TemplateArguments[0] when hasResult was true and SystemTypes.Void otherwise
                returnType: taskType,
                body: null);

            checkPostMethod.CallingConvention = CallingConventionFlags.HasThis;
            checkPostMethod.Flags |= MethodFlags.Public;

            this.checkPostBody = new StatementList();
            this.closureClass.Members.Add(this.checkPostMethod);

            if (taskType.IsGeneric)
            {
                // Assign taskParameter to originalResultLocal because
                // this field is used in a postcondition
                checkPostBody.Add(new AssignmentStatement(this.originalResultLocal, checkMethodTaskParameter));
            }
        }

        [ContractVerification(false)]
        private void AddAsyncPost(List<Ensures> asyncPostconditions)
        {
            var origBody = new Block(this.checkPostBody);
            origBody.HasLocals = true;

            var newBodyBlock = new Block(new StatementList());
            newBodyBlock.HasLocals = true;

            var methodBody = new StatementList();
            var methodBodyBlock = new Block(methodBody);
            methodBodyBlock.HasLocals = true;

            checkPostMethod.Body = methodBodyBlock;

            methodBody.Add(newBodyBlock);
            Block newExitBlock = new Block();
            methodBody.Add(newExitBlock);

            // Map closure locals to fields and initialize closure fields

            foreach (Ensures e in asyncPostconditions)
            {
                if (e == null) continue;

                this.Visit(e);

                if (this.forwarder != null)
                {
                    this.forwarder.Visit(e);
                }

                ReplaceResult repResult = new ReplaceResult(
                    this.checkPostMethod, this.originalResultLocal,
                    this.rewriter.AssemblyBeingRewritten);

                repResult.Visit(e);

                if (repResult.ContractResultWasCapturedInStaticContext)
                {
                    this.contractResultCapturedInStaticContext.Add(e.Assertion.SourceContext);
                }

                // now need to initialize closure result fields
                foreach (var target in repResult.NecessaryResultInitializationAsync(this.closureLocals))
                {
                    // note: target here 
                    methodBody.Add(new AssignmentStatement(target, this.originalResultLocal));
                }
            }

            // Emit normal postconditions

            SourceContext? lastEnsuresSourceContext = null;

            var ensuresChecks = new StatementList();

            Method contractEnsuresMethod = this.rewriter.RuntimeContracts.EnsuresMethod;

            // For generic types need to 'fix' generic type parameters that are used in the closure method. 
            // See comment to the GenericTypeMapper for more details.
            TypeParameterFixerVisitor fixer = null;

            if (!this.genericTypeMapper.IsEmpty)
            {
                fixer = new TypeParameterFixerVisitor(genericTypeMapper);
            }

            foreach (Ensures e in GetTaskResultBasedEnsures(asyncPostconditions))
            {
                // TODO: Not sure that 'break' is enough! It seems that this is possible
                // only when something is broken, because normal postconditions 
                // are using Contract.Result<T>() and this is possible only for
                // generic tasks.
                if (IsVoidTask()) break; // something is wrong in the original contract

                lastEnsuresSourceContext = e.SourceContext;

                //
                // Call Contract.RewriterEnsures
                //

                ExpressionList args = new ExpressionList();
                if (fixer != null)
                {
                    fixer.Visit(e.PostCondition);
                }

                args.Add(e.PostCondition);

                args.Add(e.UserMessage ?? Literal.Null);

                args.Add(e.SourceConditionText ?? Literal.Null);

                ensuresChecks.Add(
                    new ExpressionStatement(
                        new MethodCall(
                            new MemberBinding(null, contractEnsuresMethod),
                            args, NodeType.Call, SystemTypes.Void),
                        e.SourceContext));
            }

            this.rewriter.CleanUpCodeCoverage.VisitStatementList(ensuresChecks);

            //
            // Normal postconditions
            //

            // Wrapping normal ensures into following if statement
            // if (task.Status == TaskStatus.RanToCompletion)
            // { postcondition check }
            //
            // Implementation of this stuff is a bit tricky because if-statements
            // are inverse in the IL.
            // Basically, we need to generate following code:
            // if (!(task.Status == Task.Status.RanToCompletion))
            //   goto EndOfNormalPostcondition;
            // {postcondition check}
            // EndOfNormalPostcondition:
            // {other Code}

            // Marker for EndOfNormalPostcondition
            Block endOfNormalPostcondition = new Block();

            // Generate: if (task.Status != RanToCompletion) goto endOfNormalPostcondition;
            StatementList checkStatusStatements = CreateIfTaskResultIsEqualsTo(
                checkMethodTaskParameter, TaskStatus.RanToCompletion, endOfNormalPostcondition);

            methodBodyBlock.Statements.Add(new Block(checkStatusStatements));

            // Emit a check for __ContractsRuntime.insideContractEvaluation around Ensures
            // TODO ST: there is no sense to add recursion check in async postcondition that can be checked in different thread!
            methodBodyBlock.Statements.Add(new Block(ensuresChecks));
            // Emit a check for __ContractsRuntime.insideContractEvaluation around Ensures
            //this.rewriter.EmitRecursionGuardAroundChecks(this.checkPostMethod, methodBodyBlock, ensuresChecks);

            // Now, normal postconditions are written to the method body. 
            // We need to add endOfNormalPostcondition block as a marker.
            methodBodyBlock.Statements.Add(endOfNormalPostcondition);

            //
            // Exceptional postconditions
            //

            var exceptionalPostconditions = GetExceptionalEnsures(asyncPostconditions).ToList();
            if (exceptionalPostconditions.Count > 0)
            {
                // For exceptional postconditions we need to generate CheckException method first
                Method checkExceptionMethod = CreateCheckExceptionMethod();

                EmitCheckExceptionBody(checkExceptionMethod, exceptionalPostconditions);

                this.closureClass.Members.Add(checkExceptionMethod);

                // Then, we're using the same trick as for normal postconditions:
                // wrapping exceptional postconditions only when task.Status is TaskStatus.Faulted
                Block checkExceptionBlock = new Block(new StatementList());

                // Marker for endOfExceptionPostcondition
                Block endOfExceptionPostcondition = new Block();

                StatementList checkStatusIsException = CreateIfTaskResultIsEqualsTo(
                    checkMethodTaskParameter, TaskStatus.Faulted, endOfExceptionPostcondition);

                checkExceptionBlock.Statements.Add(new Block(checkStatusIsException));

                // Now we need to emit actuall check for exceptional postconditions

                // Emit: var ae = task.Exception;
                
                var aeLocal = new Local(aggregateExceptionType.Value);

                checkExceptionBlock.Statements.Add(
                    new AssignmentStatement(aeLocal,
                        new MethodCall(
                            new MemberBinding(checkMethodTaskParameter,
                                GetTaskProperty(checkMethodTaskParameter, "get_Exception")),
                            new ExpressionList())));

                // Emit: CheckException(ae);

                // Need to store method result somewhere, otherwise stack would be corrupted
                var checkResultLocal = new Local(SystemTypes.Boolean);

                checkExceptionBlock.Statements.Add(
                    new AssignmentStatement(checkResultLocal,
                            new MethodCall(new MemberBinding(null, 
                                checkExceptionMethod),
                            new ExpressionList(checkExceptionMethod.ThisParameter, aeLocal))));

                checkExceptionBlock.Statements.Add(endOfExceptionPostcondition);
                
                methodBody.Add(checkExceptionBlock);
            }

            // Copy original block to body statement for both: normal and exceptional postconditions.
            newBodyBlock.Statements.Add(origBody);

            Block returnBlock = CreateReturnBlock(checkMethodTaskParameter, lastEnsuresSourceContext);
            methodBody.Add(returnBlock);
        }

        /// <summary>
        /// Helper visitor class that changes all references to type parameters to appropriate once.
        /// </summary>
        private class TypeParameterFixerVisitor : StandardVisitor
        {
            private readonly GenericTypeMapper genericParametersMapping;

            public TypeParameterFixerVisitor(GenericTypeMapper genericParametersMapping)
            {
                Contract.Requires(genericParametersMapping != null);
                this.genericParametersMapping = genericParametersMapping;
            }

            public override Expression VisitAddressDereference(AddressDereference addr)
            {
                // Replacing initobj !!0 to initobj !0
                var newType = genericParametersMapping.GetClosureTypeParameterByEnclosingTypeParameter(addr.Type);
                if (newType != addr.Type)
                {
                    return new AddressDereference(addr.Address, newType, addr.Volatile, addr.Alignment, addr.SourceContext);
                }
                
                return base.VisitAddressDereference(addr);
            }

            // Literal is used when contract result compares to null: Contract.Result<T>() != null
            public override Expression VisitLiteral(Literal literal)
            {
                var origin = literal.Value as TypeNode;
                if (origin == null)
                {
                    return base.VisitLiteral(literal);
                }

                var newLiteralType = this.genericParametersMapping.GetClosureTypeParameterByEnclosingTypeParameter(origin);
                if (newLiteralType != origin)
                {
                    return new Literal(newLiteralType);
                }

                return base.VisitLiteral(literal);
            }

            public override TypeNode VisitTypeParameter(TypeNode typeParameter)
            {
                var fixedVersion = this.genericParametersMapping.GetClosureTypeParameterByEnclosingTypeParameter(typeParameter);
                if (fixedVersion != typeParameter)
                {
                    return fixedVersion;
                }

                return base.VisitTypeParameter(typeParameter);
            }

            public override TypeNode VisitTypeReference(TypeNode type)
            {
                var fixedVersion = this.genericParametersMapping.GetClosureTypeParameterByEnclosingTypeParameter(type);
                if (fixedVersion != type)
                {
                    return fixedVersion;
                }

                return base.VisitTypeReference(type);
            }

            public override TypeNode VisitTypeNode(TypeNode typeNode)
            {
                var fixedVersion = this.genericParametersMapping.GetClosureTypeParameterByEnclosingTypeParameter(typeNode);
                if (fixedVersion != typeNode)
                {
                    return fixedVersion;
                }

                return base.VisitTypeNode(typeNode);
            }
        }

        private static IEnumerable<Ensures> GetTaskResultBasedEnsures(List<Ensures> asyncPostconditions)
        {
            return asyncPostconditions.Where(post => !(post is EnsuresExceptional));
        }

        private static IEnumerable<EnsuresExceptional> GetExceptionalEnsures(List<Ensures> asyncPostconditions)
        {
            return asyncPostconditions.OfType<EnsuresExceptional>();
        }

        /// <summary>
        /// Returns TaskExtensions.Unwrap method.
        /// </summary>
        [Pure]
        private Member GetUnwrapMethod(TypeNode checkMethodTaskType)
        {
            Contract.Requires(checkMethodTaskType != null);
            Contract.Ensures(Contract.Result<Member>() != null);

            Contract.Assert(TaskExtensionsTypeNode != null, "Can't find System.Threading.Tasks.TaskExtensions type");

            var unwrapCandidates = TaskExtensionsTypeNode.GetMembersNamed(Identifier.For("Unwrap"));

            Contract.Assert(unwrapCandidates != null,
                "Can't find Unwrap method in the TaskExtensions type");

            // Should be only two methods. If that is not true, we need to change this code to reflect this!
            Contract.Assume(unwrapCandidates.Count == 2, "Should be exactly two candidate Unwrap methods.");

            // We need to find appropriate Unwrap method based on CheckMethod argument type.
            var firstMethod = (Method)unwrapCandidates[0];
            var secondMethod = (Method)unwrapCandidates[1];

            Contract.Assume(firstMethod != null && secondMethod != null);

            var genericUnwrapCandidate = firstMethod.IsGeneric ? firstMethod : secondMethod;
            var nonGenericUnwrapCandidate = firstMethod.IsGeneric ? secondMethod : firstMethod;

            if (checkMethodTaskType.IsGeneric)
            {
                // We need to "instantiate" generic first.
                // I.e. for Task<int> we need to have Unwrap(Task<Task<int>>): Task<int>
                
                // In this case we need to map back generic types.
                // CheckPost method is a non-generic method from (potentially) generic closure class.
                // In this case, if enclosing method is generic we need to map generic types back
                // and use !!0 (reference to method template arg) instead of using !0 (which is reference
                // to closure class template arg).
                var enclosingGeneritType = 
                    this.genericTypeMapper.GetEnclosingTypeParameterByClosureTypeParameter(
                        checkMethodTaskType.TemplateArguments[0]);

                return genericUnwrapCandidate.GetTemplateInstance(null, enclosingGeneritType);
            }

            return nonGenericUnwrapCandidate;
        }

        /// <summary>
        /// Factory method that creates bool CheckException(Exception e)
        /// </summary>
        [Pure]
        private Method CreateCheckExceptionMethod()
        {
            Contract.Ensures(Contract.Result<Method>() != null);

            var exnParameter = new Parameter(Identifier.For("e"), SystemTypes.Exception);
            var checkExceptionMethod =
                new Method(
                    declaringType: this.closureClass,
                    attributes: null,
                    name: CheckExceptionMethodId,
                    parameters: new ParameterList(exnParameter),
                    returnType: SystemTypes.Boolean,
                    body: new Block(new StatementList()));

            checkExceptionMethod.Body.HasLocals = true;

            checkExceptionMethod.CallingConvention = CallingConventionFlags.HasThis;
            checkExceptionMethod.Flags |= MethodFlags.Public;

            if (checkExceptionMethod.ExceptionHandlers == null)
                checkExceptionMethod.ExceptionHandlers = new ExceptionHandlerList();

            return checkExceptionMethod;
        }

        private void EmitCheckExceptionBody(Method checkExceptionMethod, List<EnsuresExceptional> exceptionalPostconditions)
        {
            Contract.Requires(checkExceptionMethod != null);

            Contract.Requires(exceptionalPostconditions != null);
            Contract.Requires(exceptionalPostconditions.Count > 0);

            // We emit the following method:
            //   bool CheckException(Exception e) {
            //     var ex = e as C1;
            //     if (ex != null) {
            //       EnsuresOnThrow(predicate)
            //     }
            //     else {
            //       var ex2 = e as AggregateException;
            //       if (ex2 != null) {
            //         ex2.Handle(CheckException);
            //       }
            //     }
            //
            //     // Method always returns true. This is by design!
            //     // We need to check all exceptions in the AggregateException
            //     // and fail in EnsuresOnThrow if the postcondition is not met.
            //     return true; // handled

            var body = checkExceptionMethod.Body.Statements;
            var returnBlock = new Block(new StatementList());

            foreach (var e in exceptionalPostconditions)
            {
                // The catchBlock contains the catchBody, and then
                // an empty block that is used in the EH.
                // TODO ST: name is confusing because there is no catch blocks in this method!
                Block catchBlock = new Block(new StatementList());

                // local is: var ex1 = e as C1;
                Local localEx = new Local(e.Type);

                body.Add(
                    new AssignmentStatement(localEx,
                        new BinaryExpression(checkExceptionMethod.Parameters[0],
                            new MemberBinding(null, e.Type),
                            NodeType.Isinst)));

                Block skipBlock = new Block();
                body.Add(new Branch(new UnaryExpression(localEx, NodeType.LogicalNot), skipBlock));
                body.Add(catchBlock);
                body.Add(skipBlock);

                // call Contract.EnsuresOnThrow
                ExpressionList args = new ExpressionList();
                args.Add(e.PostCondition);

                args.Add(e.UserMessage ?? Literal.Null);

                args.Add(e.SourceConditionText ?? Literal.Null);

                args.Add(localEx);
                var checks = new StatementList();

                checks.Add(
                    new ExpressionStatement(
                        new MethodCall(
                            new MemberBinding(null, this.rewriter.RuntimeContracts.EnsuresOnThrowMethod),
                            args,
                            NodeType.Call,
                            SystemTypes.Void),
                        e.SourceContext));

                this.rewriter.CleanUpCodeCoverage.VisitStatementList(checks);

                // TODO ST: actually I can't see this recursion guard check in the resulting IL!!
                rewriter.EmitRecursionGuardAroundChecks(checkExceptionMethod, catchBlock, checks);

                catchBlock.Statements.Add(new Branch(null, returnBlock));
            }

            // recurse on AggregateException itself
            {
                // var ae = e as AggregateException;
                // if (ae != null) {
                //   ae.Handle(this.CheckException);
                // }
                Block catchBlock = new Block(new StatementList());
                var aggregateType = aggregateExceptionType.Value;

                // var ex2 = e as AggregateException;
                Local localEx2 = new Local(aggregateType);
                body.Add(
                    new AssignmentStatement(localEx2,
                        new BinaryExpression(
                            checkExceptionMethod.Parameters[0],
                            new MemberBinding(null, aggregateType),
                            NodeType.Isinst)));

                Block skipBlock = new Block();
                body.Add(new Branch(new UnaryExpression(localEx2, NodeType.LogicalNot), skipBlock));
                body.Add(catchBlock);
                body.Add(skipBlock);

                var funcType = func2Type.Value;
                funcType = funcType.GetTemplateInstance(this.rewriter.AssemblyBeingRewritten, SystemTypes.Exception, SystemTypes.Boolean);

                var handleMethod = aggregateType.GetMethod(Identifier.For("Handle"), funcType);

                var funcLocal = new Local(funcType);
                var ldftn =
                    new UnaryExpression(
                        new MemberBinding(null, checkExceptionMethod),
                        NodeType.Ldftn,
                        CoreSystemTypes.IntPtr);

                catchBlock.Statements.Add(
                    new AssignmentStatement(funcLocal,
                        new Construct(
                            new MemberBinding(null, funcType.GetConstructor(SystemTypes.Object, SystemTypes.IntPtr)),
                            new ExpressionList(checkExceptionMethod.ThisParameter, ldftn))));

                catchBlock.Statements.Add(
                    new ExpressionStatement(new MethodCall(new MemberBinding(localEx2, handleMethod),
                        new ExpressionList(funcLocal))));
            }

            // add return true to CheckException method
            body.Add(returnBlock);
            body.Add(new Return(Literal.True));
        }

        /// <summary>
        /// Returns property for the task object.
        /// </summary>
        private static Method GetTaskProperty(Parameter taskParameter, string propertyName)
        {
            Contract.Requires(taskParameter != null);
            Contract.Ensures(Contract.Result<Method>() != null);

            // For generic task Status property defined in the base class.
            // That's why we need to check what the taskParameter type is - is it generic or not.
            // If the taskParameter is generic we need to use base type (because Task<T> : Task).

            var taskTypeWithStatusProperty = taskParameter.Type.IsGeneric
                ? taskParameter.Type.BaseType
                : taskParameter.Type;

            return taskTypeWithStatusProperty.GetMethod(Identifier.For(propertyName));

        }

        /// <summary>
        /// Method returns a list of statements that checks task status.
        /// </summary>
        private static StatementList CreateIfTaskResultIsEqualsTo(
            Parameter taskParameterToCheck, TaskStatus expectedStatus,
            Block endBlock)
        {
            Contract.Ensures(Contract.Result<StatementList>() != null);

            var result = new StatementList();

            // If-statement is slightly different in IL.
            // To get `if (condition) {statements}`
            // we need to generate:
            // if (!condition) goto endBLock; statements; endBlock:

            // This method emits a check that simplifies CheckMethod implementation.

            var statusProperty = GetTaskProperty(taskParameterToCheck, "get_Status");
            Contract.Assert(statusProperty != null, "Can't find Task.Status property");

            // Emitting: var tmpStatus = task.Status;
            var tmpStatus = new Local(statusProperty.ReturnType);
            result.Add(
                new AssignmentStatement(tmpStatus,
                    new MethodCall(new MemberBinding(taskParameterToCheck, statusProperty),
                    new ExpressionList())));

            // if (tmpStatus != expectedStatus)
            // goto endOfMethod;
            // This is an inverted form of the check: if (tmpStatus == expectedStatus) {check}
            result.Add(
                new Branch(
                    new BinaryExpression(
                        tmpStatus,
                        new Literal(expectedStatus),
                        NodeType.Ne),
                    endBlock));

            return result;
        }

        private static Block CreateReturnBlock(Parameter checkPostTaskParameter, SourceContext? lastEnsuresSourceContext)
        {
            Statement returnStatement = new Return(checkPostTaskParameter);
            if (lastEnsuresSourceContext != null)
            {
                returnStatement.SourceContext = lastEnsuresSourceContext.Value;
            }

            Block returnBlock = new Block(new StatementList(1));
            returnBlock.Statements.Add(returnStatement);

            return returnBlock;
        }

        /// <summary>
        /// Returns correct version of the ContinueWith method.
        /// </summary>
        /// <remarks>
        /// This function returns ContinueWith overload that takes TaskContinuationOptions.
        /// </remarks>
        private static Method GetContinueWithMethod(Class closureClass, TypeNode taskTemplate, TypeNode taskType)
        {
            var continueWithCandidates = taskTemplate.GetMembersNamed(Identifier.For("ContinueWith"));
            
            // Looking for an overload with TaskContinuationOptions
            const int expectedNumberOfArguments = 2;

            for (int i = 0; i < continueWithCandidates.Count; i++)
            {
                var cand = continueWithCandidates[i] as Method;
                if (cand == null) continue;

                // For non-generic version we're looking for ContinueWith(Action<Task>, TaskContinuationOptions)

                if (!taskType.IsGeneric)
                {
                    if (cand.IsGeneric) continue;

                    if (cand.ParameterCount != expectedNumberOfArguments) continue;

                    if (cand.Parameters[0].Type.GetMetadataName() != "Action`1") continue;
                    
                    if (cand.Parameters[1].Type.GetMetadataName() != "TaskContinuationOptions") continue;

                    return cand;
                }

                // For generic version we're looking for ContinueWith(Func<Task, T>, TaskContinuationOptions)
                if (!cand.IsGeneric) continue;

                if (cand.TemplateParameters.Count != 1) continue;

                if (cand.ParameterCount != expectedNumberOfArguments) continue;

                if (cand.Parameters[0].Type.GetMetadataName() != "Func`2") continue;
                if (cand.Parameters[1].Type.GetMetadataName() != "TaskContinuationOptions") continue;

                // now create instance, first of task
                var taskInstance = taskTemplate.GetTemplateInstance(
                    closureClass.DeclaringModule,
                    taskType.TemplateArguments[0]);

                // ST: some black magic is happening, but it seems it is required to get ContinueWith
                // from generic instantiated version of the task
                var candMethod = (Method)taskInstance.GetMembersNamed(Identifier.For("ContinueWith"))[i];

                // Candidate method would have following signature:
                // Task<T> ContinueWith(Task<T> t) for generic version
                return candMethod.GetTemplateInstance(null, taskType);
            }

            return null;
        }

        private static InstanceInitializer CreateConstructor(Class closureClass)
        {
            var ctor = new InstanceInitializer(closureClass, null, null, null);
            ctor.CallingConvention = CallingConventionFlags.HasThis;
            ctor.Flags |= MethodFlags.Public | MethodFlags.HideBySig;

            // Regular block that calls base class constructor
            ctor.Body = new Block(
                new StatementList(
                    new ExpressionStatement(
                        new MethodCall(new MemberBinding(ctor.ThisParameter, SystemTypes.Object.GetConstructor()),
                            new ExpressionList())),
                new Return()));
            return ctor;
        }

        private bool IsVoidTask()
        {
            return this.checkPostMethod.ReturnType == SystemTypes.Void;
        }

        class FieldRewriter : StandardVisitor
        {
            private readonly EmitAsyncClosure closure;
            private Field enclosingInstance;

            public FieldRewriter(EmitAsyncClosure closure)
            {
                this.closure = closure;
            }

            public override Expression VisitMemberBinding(MemberBinding memberBinding)
            {
                // Original postcondition could have an access to the instance state via Contract.Ensures(_state == "foo");
                // Now postcondition body resides in the different class and all references to 'this' pointer should be changed.
                // If member binding references 'this', then we need to initialize '_this' field that points to the enclosing class
                // and redirect the binding to this field.
                var thisNode = memberBinding != null ? memberBinding.TargetObject as This : null;
                if (thisNode != null && thisNode.DeclaringMethod != null && thisNode.DeclaringMethod.DeclaringType != null &&
                    // Need to change only when 'this' belongs to enclosing class instance
                    thisNode.DeclaringMethod.DeclaringType.Equals(this.closure.declaringType))
                {
                    var thisField = EnsureClosureInitialization(memberBinding.TargetObject);

                    return new MemberBinding(
                            new MemberBinding(this.closure.checkPostMethod.ThisParameter, thisField),
                            memberBinding.BoundMember);
                }

                return base.VisitMemberBinding(memberBinding);
            }

            /// <summary>
            /// Ensures that async closure has a field with enclosing class field, like public EnclosingType _this.
            /// </summary>
            private Field EnsureClosureInitialization(Expression targetObject)
            {
                if (this.enclosingInstance == null)
                {
                    var localType = this.closure.forwarder != null ? this.closure.forwarder.VisitTypeReference(targetObject.Type) : targetObject.Type;

                    var enclosedTypeField = new Field(
                        this.closure.closureClass, null, FieldFlags.Public, new Identifier("_this"), localType, null);
                    this.closure.closureClass.Members.Add(enclosedTypeField);

                    // initialize the closure field
                    var instantiatedField = Rewriter.GetMemberInstanceReference(enclosedTypeField, this.closure.closureClassInstance);
                    this.closure.ClosureInitializer.Statements.Add(
                        new AssignmentStatement(
                            new MemberBinding(this.closure.closureLocal, instantiatedField), targetObject));

                    this.enclosingInstance = enclosedTypeField;
                }

                return this.enclosingInstance;
            }
        }

        // Visitor for changing closure locals to fields
        public override Expression VisitLocal(Local local)
        {
            if (HelperMethods.IsClosureType(this.declaringType, local.Type))
            {
                MemberBinding mb;
                if (!closureLocals.TryGetValue(local, out mb))
                {
                    // TODO ST: not clear what's going on here!
                    // Clarification: this method changes access to local variables to apropriate fields.
                    // Consider following example:
                    // public async Task<int> FooAsync(int[] arguments, int length)
                    // {
                    //    Contract.Ensures(Contract.ForAll(arguments, i => i == Contract.Result<int>() && i == length)); 
                    // }
                    // In this case, CheckPost method should reference length that will become a field of the generated
                    // closure class.
                    // So this code will change all locals (like length) to appropriate fields of the async closure instance.

                    // Forwarder would be null, if enclosing method with async closure is not generic
                    var localType = forwarder != null ? forwarder.VisitTypeReference(local.Type) : local.Type;

                    var closureField = new Field(this.closureClass, null, FieldFlags.Public, local.Name, localType, null);
                    this.closureClass.Members.Add(closureField);

                    mb = new MemberBinding(this.checkPostMethod.ThisParameter, closureField);
                    closureLocals.Add(local, mb);

                    // initialize the closure field
                    var instantiatedField = Rewriter.GetMemberInstanceReference(closureField, this.closureClassInstance);
                    this.ClosureInitializer.Statements.Add(
                        new AssignmentStatement(
                            new MemberBinding(this.closureLocal, instantiatedField), local));
                }

                return mb;
            }

            return local;
        }
    }
}